#!/usr/bin/python

# Imports
import os, sys
import pygame
from data import *
from faceDetect import FaceDetect
from inputbox import Inputbox
from sql import SQL
from dialog import Dialog

# main window parameters
main_buttons = ["Login", "Take Pictures", "Create User", "Quit"]
main_window = 'Main'
main_title = 'Choose your option'

class Scene(object):
    """Base of the program. Each window gets set and drawn here.
    Backgrounds, buttons and text get drawn here.
    Each window atleast has a display screen and a window name."""
    
    def __init__(self, screen, window, window_buttons = None, users = None, title = ''):
        """Initialization of Scene where the current window is set.
        Creates a SQL, Inputbox and FaceDetect instance.
        Eventually starts the program by calling the main_loop()
        Parameters: display screen - drawing screen
                    string window - window name (choices: Main, Login, Create User, Pictures, Match, Users List) 
                    list window_buttons - list of button names
                    list users - list of users where each user contains an user_id, title, first and last name in this order
                    string title - title of the window"""
        print 'Initializing Scene'
        self.buttons = []
        self.keepGoing = True
        
        self.set_scene(screen, window, window_buttons, users, title)
        
        self.sql = SQL()
        self.sql.connect_db()
        self.box = Inputbox(self.screen, self.sql, self)
        self.FD = FaceDetect(self.screen, self.sql, self)
        
        self.main_loop()
    
    def main_loop(self):        
        """Starting the loop of the program
        Quits the program when you're in the main window and hit escape, else it will just return to it.
        This loop checks which window should be drawn by looking at self.window.
        Each window has it's own buttons and track_buttons. Track_buttons checks which button is pressed."""
        while self.keepGoing:
            events = pygame.event.get()
            for event in events:
                if event.type == pygame.QUIT:
                    self.sql.close_db()
                    sys.exit()
                    print 'Quitting Facerecognition'
                elif event.type == pygame.KEYUP:
                    if event.key == pygame.K_ESCAPE:
                        if self.window == 'Main':
                            print 'Quitting Facerecognition'
                            self.sql.close_db()
                            sys.exit()
                        else:
                            self.main_window()
                elif event.type == pygame.MOUSEBUTTONUP:
                    if self.window == 'Main':
                        self.track_main_buttons()
                    elif self.window == 'Login':
                        self.track_login_buttons()
                    elif self.window == 'Create User':
                        if self.track_create_buttons(self):
                            exception = False
                            log_name = ''
                            if " " in self.inputbox.capitalize():
                                first_name, last_name = self.inputbox.split(" ")
                                log_name = '%s_%s' % (first_name.capitalize(), last_name.capitalize())
                            else:
                                log_name = '%s_Achternaam' % (self.inputbox.capitalize())
                            
                            if self.sql.user_exists(login_name = log_name):
                                exception = True
                            else:
                                self.sql.create_user(self.inputbox.capitalize())
                                
                            if not exception:
                                user = self.sql.get_user(user_id = self.sql.get_last_inserted_user_id())
                                dialog = Dialog(self.screen, 'Create User', \
                                       ['User named %s %s created.' % (user[2], user[3]), 'User Identification Code: %s' % (user[0])],  \
                                        'Shoot Picture', self)
                                dialog.set_user_id(user[0])
                                dialog.main()
                            else:
                                user = self.sql.get_user(login_name = log_name)
                                dialog = Dialog(self.screen, 'Error: Create User', \
                                       ['User named %s %s already exists.' % (user[1], user[2]), 'User Identification Code: %s' % (user[0])],  \
                                        'Shoot Picture', self)
                                dialog.set_user_id(user[0])
                                dialog.main()
                    elif self.window == 'Pictures':
                        self.track_picture_buttons(self.inputbox, self.img, self.box)
                    elif self.window == 'Match':
                        self.track_match_buttons()
                    elif self.window == 'Users List':
                        self.track_user_list_buttons()
            
            self.screen.fill((0, 0, 0))
            if self.window == 'Main':
                self.draw_background("main_background.png")
                self.draw_window_buttons(vertical = True)
            elif self.window == 'Login':
                self.draw_background("background.png")
                self.draw_window_buttons(vertical = False)
                self.img = self.FD.login()
            elif self.window == 'Create User':
                self.draw_background("background.png")
                self.draw_window_buttons(vertical = False)
                self.draw_text('Only one First and Last name is possible. (Example: Jan Jansen)', 
                               28, (255, 255, 255), 'CENTER', 325)
                self.box.set_question('Name')
                self.inputbox = self.box.ask(events)
            elif self.window == 'Pictures':
                self.draw_background("background.png")
                self.draw_window_buttons(vertical = False)
                self.img = self.FD.draw_cam()
                self.draw_text('Enter your user identification code.', 
                               28, (255, 255, 255), 'CENTER', 360)
                self.box.set_question('User ID')
                self.inputbox = self.box.ask(events, 1.5)
            elif self.window == 'Match':
                self.draw_background("main_background.png")
                self.draw_matches()
                self.draw_window_buttons(vertical = False)
            elif self.window == 'Users List':
                self.draw_background("main_background.png")
                self.draw_matches()
            self.draw_buttons()
            self.draw_text(self.title, 40, (255, 255, 255), 'CENTER', 55)
            pygame.display.flip()
    
    def draw_background(self, background):
        """Draws the background picture on the screen.
        Parameters: string background - name of the background image"""
        background = load_image(os.path.join(DATADIR, background))
        self.screen.blit(background, (0, 0))
    
    def track_main_buttons(self):
        """Tracks all buttons in the main window, performing an action depending on the button pressed."""
        for b in self.window_buttons:
            if mouse_collision(b[1]):
                print '[%s] button was pressed with %s' % (b[0], b[1])
                if b[0] == 'Login':
                    self.login_window()
                if b[0] == 'Take Pictures':
                    self.take_picture_window()
                if b[0] == 'Create User':
                    self.create_user_window()
                if b[0] == 'Quit':
                    sys.exit()
    
    def track_login_buttons(self):
        """Tracks all buttons in the login window, performing an action depending on the button pressed.
        Try Again - Resets the countdown, starting over again.
        Back - Returns to the main window"""
        for b in self.window_buttons:
            if mouse_collision(b[1]):
                print '[%s] button was pressed with %s' % (b[0], b[1])
                if b[0] == 'Try Again':
                    self.FD.reset_countdown()
                if b[0] == 'Back':
                    self.main_window()

    def track_create_buttons(self, box):
        """Tracks all buttons in the create user window, performing an action depending on the button pressed.
        Create - Returns True which simply determines if the button is pressed.
        Back - Cleans the inputbox and returns to the main window
        Parameters: Inputbox box - instance of Inputbox()"""
        for b in self.window_buttons:
            if mouse_collision(b[1]):
                print '[%s] button was pressed with %s' % (b[0], b[1])
                if b[0] == 'Create':
                    return True
                if b[0] == 'Back':
                    self.box.reset()
                    self.main_window()
    
    def track_picture_buttons(self, user_id, img, box):
        """Tracks all buttons in the take picture window, performing an action depending on the button pressed.
        Shoot - Takes a picture, saves it in 'capture/' and in the database, linking it to the user_id entered in the inputbox.
        Back - Cleans the inputbox and returns to the main window
        Parameters: int user_id - user identification code
                    IplImage img - image of frame of the (web)cam
                    Inputbox box - instance of Inputbox()"""
        for b in self.window_buttons:
            if mouse_collision(b[1]):
                print '[%s] button was pressed with %s' % (b[0], b[1])
                if b[0] == 'Shoot':
                    if self.sql.user_exists(user_id = user_id):
                        self.FD.set_user_id(int(user_id))
                        self.FD.save_tmp_image(img)
                        for filename in self.FD.captures:
                            self.FD.save_diffs(filename)
                            self.FD.captures = []
                    else:
                        box.existing_user = False
                if b[0] == 'Back':
                    self.box.reset()
                    self.main_window()
    
    def track_match_buttons(self):
        """Tracks all buttons in the match window, performing an action depending on the button pressed.
        Try Again - Reloads the login window.
        Main menu - Return to the main window.
        Users List - Shows the users list window and sets self.match_found = True, determing the window title that will be shown.
        Reload - Refines the matches.
        The matches (names) shown are buttons aswell which saves and links the captured picture(s) to the name when pressed. """
        for b in self.window_buttons:
            if mouse_collision(b[1]):
                print '[%s] button was pressed with %s' % (b[0], b[1])
                if b[0] == 'Try Again':
                    self.login_window()
                if b[0] == 'Main menu':
                    self.main_window()
                if b[0] == 'Users List':
                    self.user_list_window(True)
        for t in self.users:
            if mouse_collision(t[1]):
                print '[%s] button was pressed with %s' % (t[0], t[1])
                self.learn_system(t)
    
    def track_user_list_buttons(self):
        """Tracks all buttons in the user list window, performing an action depending on the button pressed.
        Back - Return to the main window.
        The names shown are buttons aswell which saves and links the captured picture(s) to the name when pressed. """
        for b in self.buttons:
            if mouse_collision(b[1]):
                print '[%s] button was pressed with %s' % (b[0], b[1])
                if b[0] == 'Back':
                    self.main_window()
        for t in self.users:
            if mouse_collision(t[1]):
                print '[%s] button was pressed with %s' % (t[0], t[1])
                self.learn_system(t)
    
    def learn_system(self, user):
        """Saves and links the captured picture(s) to the name when pressed. 
        Saving the picture in 'capture/' and linking it to the user through user_id by creating an user_face in the database.
        Parameters: list user: list of user parameters which should contain user_id, title, first and last name in this order."""
        user_id, title, name = user[0][0], user[0][1], '%s %s ' % (user[0][2], user[0][3])
        print 'Selected user:', 'ID = %s, Title = %s, Name = %s' % (user_id, title, name)
        self.FD.set_user_id(user_id)
        for capture in self.FD.captures:
            self.FD.save_diffs(capture)
        self.main_window()
    
    def set_scene(self, screen, window, window_buttons = None, users = None, title = ''):
        """Sets the current scene.
        Parameters: display screen - drawing screen
                    string window - window name (choices: Main, Login, Create User, Pictures, Match, Users List) 
                    list window_buttons - list of button names
                    list text - list users (where each user contains an user_id, title, first and last name in this order)
                    string title - title of the window"""
        self.screen = screen
        self.window = window
        self.title = title
        self.window_buttons = []
        self.users = []
        self.buttons = []
        b = load_image(os.path.join(DATADIR, "button_blue.png"))
        b_ro = load_image(os.path.join(DATADIR, "button_black.png"))
        if window_buttons:
            for button in window_buttons:
                self.window_buttons.append([button, pygame.Rect(0, 0, 0, 0),  b, b_ro])
        if users:
            for line in users:
                self.users.append([line, pygame.Rect(0, 0, 0, 0)])
    
    def create_button(self, name, pos, b, b_ro):
        """Adds a button to the buttons list. This function can be used for different buttons on a specific locations.
        Parameters: string name - name of the button
                    topleft pos - position of the button (x, y)
                    string b - name of the button image
                    string b_ro - name of the button image when selected"""
        b = load_image(os.path.join(DATADIR, b))
        b_ro = load_image(os.path.join(DATADIR, b_ro))
        rect = pygame.Rect(b.get_rect())
        rect.x, rect.y = pos
        self.buttons.append([name, rect, b, b_ro])
    
    def draw_text(self, text, fontsize, color, x, y):
        """Draws a text line
        Parameters: string text - Text string
                    int fontsize - The size of the font
                    Color color - color of the text
                    string x - x position of the text (LEFT, CENTER or RIGHT)
                    int y - y position of the text"""
        font = pygame.font.Font(None, fontsize)
        ren = font.render(text, 1, color)
        if x == 'LEFT':
            x = 25
        elif x == 'CENTER':
            x = (self.screen.get_width() / 2) - (ren.get_width() / 2)
        elif x == 'RIGHT':
            x = self.screen.get_width() - ren.get_width() - 25
        self.screen.blit(ren, (x, y))
    
    def draw_matches(self):
        """Draws the match (user) buttons showing their Title, First and Last name."""
        font = pygame.font.Font(None, 36)
        width = 0
        height = 0
        i = 1
        j = 0
        widest = 0 
        color = (255, 255, 255)
        
        if not self.match_found:
            # Replaces title
            self.title = ''
            self.draw_text('No face detected. Please choose your name or try again.', 
                               28, (255, 255, 255), 'CENTER', 55)
        
        for t in self.users:
            if mouse_collision(t[1]):
                color = (0, 255, 0)
            else:
                color = (255, 255, 255)
            
            title_ren = font.render(t[0][1], 1, color)
            name_ren = font.render(t[0][2] + ' ' + t[0][3], 1, color)
            
            if name_ren.get_width() > widest:
                widest = name_ren.get_width()
            
            if title_ren.get_width() >= name_ren.get_width():
                width = title_ren.get_width()
            else:
                width = name_ren.get_width()
            height = (title_ren.get_height() + name_ren.get_height())
            x = 25
            y = (title_ren.get_height() + name_ren.get_height())
            
            if (i % 9 == 0):
                j += 1
                i = 1
            
            x_point = x + (j * (widest + 10))
            y_point = (i * y) + 43 + (i * 10)
            t[1] = pygame.Rect(x_point, y_point, width, height)
            self.screen.blit(title_ren, (x_point, y_point))
            self.screen.blit(name_ren, (x_point, y_point + title_ren.get_height()))
            i += 1
    
    def draw_buttons(self):
        """Draws all buttons that were created by create_button()."""
        for b in self.buttons:
            self.screen.blit(b[2], b[1])
            if mouse_collision(b[1]):
                self.screen.blit(b[3], b[1])
            
    def draw_window_buttons(self, vertical = True):
        """Draws all window buttons specified for each window.
        Parameters: bool vertical - if True then the buttons are drawn vertically, if False then the buttons are drawn horizontally."""
        font = pygame.font.Font(None, 40)
        h = 1
        v = 1
        width = 0
        height = 0
        
        for b in self.window_buttons:
            total_width = (len(self.window_buttons) * (b[2].get_width())) + ((len(self.window_buttons) - 1)  * 20)
            if vertical:
                width = ((self.screen.get_width()/ 2) - (b[2].get_width() / 2))
                height = (b[2].get_height() + 40)
                v = 1
            else:
                width = (self.screen.get_width() / 2) - (total_width / 2)
                height = (self.screen.get_height() / 1.25) - (b[2].get_height() / 2)
                h = 1
            b[1] = b[2].get_rect()
            b[1].x, b[1].y = v*(width+10) + ((v-1)*b[1].width - (v-1)*width),  h*height
            self.screen.blit(b[2], b[1])
            if mouse_collision(b[1]):
                self.screen.blit(b[3], b[1])
            ren = font.render(b[0], 1, (255, 255, 255))
            self.screen.blit(ren, ((b[1].width / 2) - (ren.get_width() / 2) + (v * (width + 10) + ((v-1)*b[1].width - (v-1)*width)), (h * height) + ren.get_height()))
            h += 1
            v += 1
    
    def main_window(self):
        """Sets the scene to main window"""
        print 'Opening main window'
        self.set_scene(self.screen, main_window, main_buttons, title = main_title)
        
    def login_window(self):
        """Sets the scene to login window"""
        print 'Opening login window'
        self.match_found = True
        self.FD.captures = []
        self.FD.reset_countdown()
        self.set_scene(self.screen, 'Login', ['Try Again','Back'])
    
    def take_picture_window(self):
        """Sets the scene to take picture window"""
        print 'Opening picture window'
        self.set_scene(self.screen, 'Pictures', ['Shoot','Back'])
    
    def create_user_window(self):
        """Sets the scene to create user window"""
        print 'Opening create user window'
        self.set_scene(self.screen, 'Create User', ['Create', 'Back'])
    
    def match_window(self, users):
        """Sets the scene to match window
        Parameters: list users - list of users where each user contains an user_id, title, first and last name in this order"""
        print 'Opening match window'
        self.set_scene(self.screen, 'Match', ['Try Again', 'Main menu', 'Users List'], users, "Found Matches")
    
    def user_list_window(self, match_found):
        """Sets the scene to user list window
        Parameters: bool match_found - Determines the title of this window."""
        print 'Opening user list window'
        self.match_found = match_found
        self.set_scene(self.screen, 'Users List', None, self.sql.get_user_list(), 'User List')
        self.create_button('Back', (610, 10), 'back.png', 'back_ro.png')

def main(screen):
    """Creates a Scene instance which sets the scene to main window."""
    Scene(screen, main_window, main_buttons, title = main_title)
